import torch
from torch import nn
from torchsummary import summary
import torch.nn.functional as F
import os

from dataset import GetDataLoader
'''
AlexNet的网络是227x227的大小
'''

class AlexNet(nn.Module):
    # 模型参数保存位置
    PTH_PATH = os.path.join('./model_files','best-model.pth')
    # 输入图片的通道
    IMG_CHANNELS = GetDataLoader.IMG_CHANNELS
    # 数据集图片的分类数
    IMG_CLASSES = len(GetDataLoader().IMG_CLASSES)
    def __init_params(self):
        ## 初始化参数
        for m in self.modules():
            # 卷积层参数的初始化，w和b参数
            if isinstance(m, nn.Conv2d):
                ## 对于卷积，使用凯明初始化（对w初始化）
                nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                ## 对b参数进行初始
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            # 全连接成参数的初始化
            elif isinstance(m, nn.Linear):
                ## 默认使用正态分布
                nn.init.normal_(m.weight, 0, 0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def __init__(self):
        super(AlexNet,self).__init__()
        # 参数初始化
        self.__init_params()
        # 网络结构
        self.sequent = nn.Sequential(
            # 一次卷积使用一次激活  (该层的参数:卷积核h*w,11*11)
            nn.Conv2d(in_channels=self.IMG_CHANNELS,out_channels=96,kernel_size=11,stride=4),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=3,stride=2),

            # 一次卷积使用一次激活  (该层的参数:卷积核h*w,5*5)
            nn.Conv2d(in_channels=96,out_channels=256,kernel_size=5,padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2,stride=2),

            # 一次卷积使用一次激活  (该层的参数:卷积核h*w,3*3)
            nn.Conv2d(in_channels=256,out_channels=384,kernel_size=3,padding=1),
            nn.ReLU(),

            # 一次卷积使用一次激活  (该层的参数:卷积核h*w,3*3)
            nn.Conv2d(in_channels=384,out_channels=384,kernel_size=3,padding=1),
            nn.RReLU(),

            # 一次卷积使用一次激活  (该层的参数:卷积核h*w,3*3)
            nn.Conv2d(in_channels=384,out_channels=256,kernel_size=3,padding=1),
            nn.ReLU(),

            nn.MaxPool2d(kernel_size=3,stride=2),

            # 展平时是: 256*6*6 其中256是通道数
            nn.Flatten(),

            # 全连接层1,out_features代表的是神经元个数  (此层的参数个数=256*6*6*4096)
            nn.Linear(in_features=256*6*6,out_features=4096),
            nn.ReLU(),
            # 随机删除神经元
            nn.Dropout2d(p=0.5),

            # 全连接层2 (此层的参数个数=4096*4096)
            nn.Linear(in_features=4096,out_features=4096),
            nn.ReLU(),
            # 随机删除神经元
            nn.Dropout2d(p=0.5),

            # 最后的全连接层,不做激活,不dropout (此层的参数个数=4096*10)
            nn.Linear(in_features=4096,out_features=self.IMG_CLASSES),

            ## 最后一个全连接网络的out_features就是分类的个数
        )

    def forward(self,input):
        '''
        :param input: 输入的数据是(1,227,227)  1是通道, 在dataset中需要设置好
        :return:
        '''
        output = self.sequent(input)
        return output

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = AlexNet()
    model = model.to(device)
    print(summary(model,(3,227,227)))